from typing import Any

from vnpy.trader.constant import Direction

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)


class BollDemoMovingAtrSlStrategy(CtaTemplate):
    """
    布林线demo策略-移动止损-atr
    """
    author = "Cjie Wong"

    # 计算布林线的窗口大小
    boll_window = 18
    # 中轨标准差的倍数
    boll_dev = 3.4
    # 每次下单的数目
    fixed_size = 1
    # 计算atr指标的窗口大小
    atr_window = 20
    # 计算止损值的atr乘数
    atr_multiplier = 2
    # 固定的止盈值（take profit）
    fixed_tp = 20

    # 布林线上线值
    boll_up = 0
    # 布林线下线值
    boll_down = 0
    # 布林线中线值
    boll_mid = 0

    # atr的值
    atr_value = 0
    # 内部缓存的最高价格
    intra_trade_high = 0
    # 多头止损价
    long_sl = 0
    # 内部缓存的最低价格
    intra_trade_low = 0
    # 空头止损价
    short_sl = 0

    # 多头入场价
    long_entry = 0
    # 多头止盈价
    long_tp = 0
    # 空头入场价
    short_entry = 0
    # 空头止盈价
    short_tp = 0


    parameters = [
        "boll_window", "boll_dev", "fixed_size",
        "atr_window", "atr_multiplier",
        "fixed_tp"
    ]

    variables = [
        "boll_up", "boll_down", "boll_mid",
        "intra_trade_high", "long_sl",
        "intra_trade_low", "short_sl",
        "atr_value", "long_entry", "long_tp",
        "short_entry", "short_tp"
    ]

    def __init__(self, cta_engine: Any, strategy_name: str, vt_symbol: str, setting: dict) -> None:
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        self.bg = BarGenerator(self.on_bar, 15, self.on_15min_bar)
        self.am = ArrayManager()

    def on_init(self) -> None:
        """
        Callback when strategy is inited
        """
        self.write_log("策略初始化,载入10天的历史数据")
        self.load_bar(10)

    def on_start(self) -> None:
        """
        Callback when strategy is started
        """
        self.write_log("策略启动")

    def on_stop(self) -> None:
        """
        Callback when strategy is stopped
        """
        self.write_log("策略停止")

    def on_tick(self, tick: TickData) -> None:
        """
        Callback of new tick data update.
        """
        self.bg.update_tick(tick)

    def on_bar(self, bar: BarData) -> None:
        """
        Callback of new bar data update.
        """
        self.bg.update_bar(bar)

    def on_15min_bar(self, bar: BarData):
        """"""
        self.cancel_all()
        am = self.am
        am.update_bar(bar)
        if not am.inited:
            return
        # 计算此刻布林线的上线和下线对应的数值
        self.boll_up, self.boll_down = am.boll(self.boll_window, self.boll_dev)
        # 计算此刻布林线的中线对应的数值，即简单移动平均线的值
        self.boll_mid = am.sma(self.boll_window)
        # 计算atr的值
        self.atr_value = am.atr(self.atr_window)

        # 如果没有持仓
        if (self.pos == 0):
            # 以布林线上线的值作为价格，发起一个买入的停止单（若此刻市场价高于这个值，会立即触发以市价成交）
            self.buy(self.boll_up, self.fixed_size, True)
            # 以布林线下线的值作为价格，发起一个卖出的停止单（若此刻市场价低于这个值，会立即触发以市价成交）
            self.short(self.boll_down, self.fixed_size, True)

            # 没有持仓的时候，需要将当前K线的最高价和最低价缓存起来
            self.intra_trade_high = bar.high_price
            self.intra_trade_low = bar.low_price

            # 没有持仓的时候，初始化long_entry,long_tp,short_entry,short_tp
            self.long_entry = 0
            self.long_tp = 0
            self.short_entry = 0
            self.short_tp = 0

        # 如果持有多头仓位
        elif self.pos > 0:
            # 持有多头仓位的时候，需要更新下缓存的内部最高价。
            self.intra_trade_high = max(self.intra_trade_high, bar.high_price)
            self.intra_trade_low = bar.low_price

            # 多头止损价=内部最高价-atr的值*atr的止损乘数
            self.long_sl = self.intra_trade_high - self.atr_value * self.atr_multiplier
            # 为了避免发多个停止单，需要将止损停止单和出场停止单合二唯一，取布林线中线值和多头止损价的较大值
            self.long_sl = max(self.boll_mid, self.long_sl)
            # 创建一个卖出停止单，价格只要跌破long_sl就触发市价卖出
            self.sell(self.long_sl, abs(self.pos), True)

            # 下一个止盈单 stop profit !!!!止盈绝对不能用停止单
            if self.long_tp:
                self.sell(self.long_tp, abs(self.pos))

        # 如果持有空头仓位
        elif self.pos < 0:
            # 持有空头仓位时候，需要更新下缓存中的内部最低价
            self.intra_trade_low = min(self.intra_trade_low, bar.low_price)
            self.intra_trade_high = bar.high_price

            # 空头止损价格 = 内部最低价 + art的值*art的止损乘数
            self.short_sl = self.intra_trade_low + self.atr_value * self.atr_multiplier
            # 为了避免发多个停止单，需要将止损停止单和出场停止单合二唯一，取布林线中线值和空头止损价的较小值
            self.short_sl = min(self.boll_mid, self.short_sl)
            # 创建一个停止单，价格只要突破short_sl就会触发市价买入平仓
            self.cover(self.short_sl, abs(self.pos), True)

            if self.short_tp:
                self.cover(self.short_tp, abs(self.pos))

        self.put_event()

    def on_trade(self, trade: TradeData) -> None:
        """
        Callback of new trade data update.
        """
        if self.pos != 0:
            if trade.direction == Direction.LONG:
                self.long_entry = trade.price
                self.long_tp = self.long_entry + self.fixed_tp
            else:
                self.short_entry = trade.price
                self.short_tp = self.short_entry - self.fixed_tp

        self.put_event()

    def on_order(self, order: OrderData) -> None:
        """
        Callback of new order data update
        """
        pass

    def on_stop_order(self, stop_order: StopOrder) -> None:
        """
        Callback of stop order update
        """
        pass
